import tensorflow as tf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Bidirectional
import mplfinance as mpf
import random
import tushare as ts
import math

def set_seeds():
    random.seed(1)
    np.random.seed(1)
    tf.random.set_seed(1)


def get_stock_data(token, ts_code, start_date, end_date, fields):
    """
    从tushare获取指定股票代码、日期范围以及字段的股票数据

    参数:
    token (str): 在tushare平台注册获取的token
    ts_code (str): 股票代码，如 "601939.SH"
    start_date (str): 开始日期，格式为 "YYYYMMDD"
    end_date (str): 结束日期，格式为 "YYYYMMDD"
    fields (list): 要获取的股票数据字段列表，如 ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    返回:
    pd.DataFrame: 获取到的股票数据，如果获取失败返回None
    """
    try:
        pro = ts.pro_api('572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46')
        data = pro.daily({
            "ts_code": '002594.SZ',
            "trade_date": "",
            "start_date": '20241120',
            "end_date": '20241220',
            "fields": fields,
            "offset": "",
            "limit": ""
        })
        return data
    except Exception as e:
        print(f"获取股票数据时出现错误: {e}")
        return None


def normalize_data(data):
    """
    对输入数据进行归一化处理

    参数:
    data (numpy.ndarray): 输入的股票价格数据，形状为 (样本数量, 1)

    返回:
    tuple: 包含归一化后的数据以及对应的归一化器对象，方便后续反归一化使用
    """
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data)
    return scaled_data, scaler


def split_data(data, n_timestamp):
    """
    将输入的时间序列数据按照指定的时间步长进行分割，用于训练模型

    参数:
    data (numpy.ndarray): 归一化后的时间序列数据，形状为 (样本数量, 1)
    n_timestamp (int): 时间步长，即每个训练样本包含的时间序列长度

    返回:
    tuple: 分割后的输入数据 (X) 和对应的目标输出数据 (y)，形状分别为 (样本数量 - 时间步长, 时间步长, 1) 和 (样本数量 - 时间步长, 1)
    """
    X, y = [], []
    for i in range(len(data) - n_timestamp):
        X.append(data[i:i + n_timestamp])
        y.append(data[i + n_timestamp])
    return np.array(X), np.array(y)


def build_model(model_type, input_shape, output_shape):
    """
    根据指定的模型类型构建相应结构的深度学习模型（LSTM相关）

    参数:
    model_type (int): 模型类型选择，1表示单层LSTM，2表示多层LSTM，3表示双向LSTM
    input_shape (tuple): 输入数据的形状，通常为 (时间步长, 特征数量)
    output_shape (int): 输出数据的维度，即预测的目标数量（这里通常为1，预测股价）

    返回:
    tf.keras.models.Sequential: 构建好的深度学习模型
    """
    model = Sequential()
    if model_type == 1:
        # 单层 LSTM
        model.add(LSTM(units=50, activation='relu',
                       input_shape=input_shape))
        model.add(Dense(units=output_shape))
    elif model_type == 2:
        # 多层 LSTM
        model.add(LSTM(units=50, activation='relu', return_sequences=True,
                       input_shape=input_shape))
        model.add(LSTM(units=50, activation='relu'))
        model.add(Dense(1))
    elif model_type == 3:
        # 双向 LSTM
        model.add(Bidirectional(LSTM(50, activation='relu'),
                                input_shape=input_shape))
        model.add(Dense(1))
    model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
                  loss='mean_squared_error')
    return model


if __name__ == "__main__":
    set_seeds()
    # 股票相关参数
    token = '572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46'  # 替换为真实token
    ts_code = "002594.SZ"  # 建设银行股票代码，可替换为其他感兴趣的股票代码
    start_date = "20100101"
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 深度学习模型相关参数
    n_timestamp = 5  # 时间戳窗口大小
    n_epochs = 30  # 训练轮数
    model_type = 1  # 1: 单层 LSTM; 2: 多层 LSTM; 3: 双向 LSTM

    # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)
    if stock_data is None:
        print("无法获取股票数据，程序退出。")
        exit(1)
    # 将数据存储为CSV文件，方便后续查看和使用
    stock_data.to_csv("stock_{}.csv".format(ts_code), index=False)

    # 读取数据并进行预处理
    mydata = pd.read_csv("stock_{}.csv".format(ts_code), parse_dates=["trade_date"], index_col="trade_date")[["open", "high", "low", "close", "vol"]]
    mydata = mydata.iloc[::-1]  # 翻转数据，可根据实际需求决定是否需要这步操作

    # 获取股价数据部分（这里可根据需求调整取数逻辑，比如取不同时间范围等）
    df3 = mydata.reset_index().loc[30:, ['trade_date', 'open', 'high', 'low', 'close', 'vol']]
    print("生成df3后的基本信息：")
    print(df3.info())
    print("生成df3后的列信息：", df3.columns)
    print("生成df3后的数据形状:", df3.shape)
    df3 = df3.dropna(how='any').reset_index(drop=True)
    print("去除缺失值后df3的基本信息：")
    print(df3.info())
    print("去除缺失值后df3的数据形状:", df3.shape)
    df3 = df3.sort_values(by='trade_date', ascending=True)
    print("排序后df3的基本信息：")
    print(df3.info())
    print("排序后df3的数据形状:", df3.shape)
    df3.set_index('trade_date', inplace=True)
    print("重新设置索引后df3的索引类型:", df3.index.dtype)
    # 重命名列名，将'vol'修改为'volume'，使其符合mplfinance要求
    df3.rename(columns={'vol': 'volume'}, inplace=True)
    print("重命名列名后df3的列信息:", df3.columns)

    import matplotlib.pyplot as plt

    plt.rcParams['font.sans-serif'] = ['SimSun']
    plt.rcParams['axes.unicode_minus'] = False
    # 可视化蜡烛图及均线
    import matplotlib.pyplot as plt
    import mplfinance as mpf
    plt.rc('font', family='SimHei')
    plt.rcParams['axes.unicode_minus'] = False
# 数据处理部分（假设已经有处理好的数据df3）

# 绘图操作
    mpf.plot(df3,
         type='candle',
         style='yahoo',
         title='股票价格走势及均线',
         ylabel='价格',
         ylabel_lower='成交量',
         volume=True,
         figsize=(12, 8),
         savefig='stock_candle.png',
         addplot=[mpf.make_addplot(df3['close'].rolling(window=5).mean()),
                  mpf.make_addplot(df3['close'].rolling(window=10).mean()),
                  mpf.make_addplot(df3['close'].rolling(window=30).mean())],
         warn_too_much_data=10000
         )
         
    # 准备用于深度学习的数据
    training_set = df3.close.values.reshape(-1, 1)
    test_set = df3.close.values[-300:].reshape(-1, 1)

    # 数据归一化
    training_set_scaled, scaler = normalize_data(training_set)
    testing_set_scaled = scaler.transform(test_set)

    # 数据分割
    X_train, y_train = split_data(training_set_scaled, n_timestamp)
    X_test, y_test = split_data(testing_set_scaled, n_timestamp)

    # 调整数据形状，符合LSTM模型输入要求（增加维度用于表示样本数量、时间步长和特征数量）
    X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
    X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

   # 构建模型
    model = build_model(model_type, (X_train.shape[1], 1), 1)
    model.summary()
    # 训练模型
    history = model.fit(X_train, y_train,
                        batch_size=64,
                        epochs=n_epochs,
                        validation_data=(X_test, y_test),
                        validation_freq=1,
                        verbose=1)

    # 绘制训练和验证损失曲线
    plt.plot(history.history['loss'], label='训练损失')
    plt.plot(history.history['val_loss'], label='验证损失')
    plt.title("模型训练和验证损失曲线")
    plt.xlabel("轮数")
    plt.ylabel("损失值")
    plt.legend()
    plt.show()

    # 模型预测
    predicted_stock_price = model.predict(X_test)
    predicted_stock_price = scaler.inverse_transform(predicted_stock_price)
    real_stock_price = scaler.inverse_transform(y_test)

    # 可视化真实价格和预测价格对比
    plt.figure(figsize=(12, 5))
    plt.plot(real_stock_price, label='真实股价')
    plt.plot(predicted_stock_price, label='预测股价')
    plt.title("真实股价与预测股价对比")
    plt.xlabel("时间")
    plt.ylabel("股价")
    plt.legend()
    plt.show()